//
// Created by Tony on 2024/5/20.
//

#include "StmtParser.h"

#include "../common/KeyWords.h"
#include "../common/Utils.h"
#include "../lexer/token/NameToken.h"
#include "ast/Ast.h"
#include "ast/list/stmt/BlockStmt.h"
#include "ast/list/stmt/IfStmt.h"
#include "ast/list/stmt/ImportStmt.h"
#include "ast/list/stmt/ReturnStmt.h"
#include "ast/list/stmt/SimpleStmt.h"
#include "ast/list/stmt/VarDefStmt.h"
#include <memory>

namespace Parser {

using std::make_shared;

shared_ptr<Ast> StmtParser::parse() { return stmt(); }

/**
 * simple_stmt := expr ";"
 */
shared_ptr<SimpleStmt> StmtParser::simpleStmt() {
    auto expr = this->expr();
    CHECK_NULL_RETURN(expr, nullptr);
    if (!skipPunct(";")) {
        return nullptr;
    }
    return make_shared<SimpleStmt>(vector<shared_ptr<Ast>>{expr});
}

/**
 * var_def_stmt := "var" NAME ( [ type_tag ] | [ "=" expr ] ) ";"
 */
shared_ptr<VarDefStmt> StmtParser::varDefStmt() {
    CHECK_FALSE_RETURN(skipName(KeyWords::VAR), nullptr);

    auto stmt = make_shared<VarDefStmt>();

    auto name = this->name();
    CHECK_NULL_RETURN(name, nullptr);
    stmt->setName(name);

    if (isPunct(tokenQueue.peek(), ":")) {
        auto typeTag = this->typeTag();
        CHECK_NULL_RETURN(typeTag, nullptr);
        stmt->setTypeTag(typeTag);
    }

    if (isPunct(tokenQueue.peek(), "=")) {
        tokenQueue.poll();
        auto initStmt = this->expr();
        CHECK_NULL_RETURN(initStmt, nullptr);
        stmt->setInitExpr(initStmt);
    }

    if (stmt->getTypeTag() == nullptr && stmt->getInitExpr() == nullptr) {
        setErrMsg("expect type tag or initialize statement for var define at %s",
                  name->position().startStr().c_str());
        return nullptr;
    }

    CHECK_FALSE_RETURN(skipPunct(";"), nullptr);
    return stmt;
}

/**
 * if_stmt := "if" "(" expr ")" block { "else" (block | if_stmt) }
 */
shared_ptr<IfStmt> StmtParser::ifStmt() {
    CHECK_FALSE_RETURN(skipName(KeyWords::IF), nullptr);
    CHECK_FALSE_RETURN(skipPunct("("), nullptr);

    auto ifStmt = make_shared<IfStmt>();

    auto conditionExpr = expr();
    CHECK_NULL_RETURN(conditionExpr, nullptr);
    ifStmt->setCondition(conditionExpr);

    CHECK_FALSE_RETURN(skipPunct(")"), nullptr);

    auto ifBlock = this->block();
    CHECK_NULL_RETURN(ifBlock, nullptr);
    ifStmt->setIfBlock(ifBlock);

    if (isName(tokenQueue.peek(), KeyWords::ELSE)) {
        tokenQueue.poll();

        shared_ptr<Ast> elseBlock;
        if (isName(tokenQueue.peek(), KeyWords::IF)) {
            elseBlock = this->ifStmt();
        } else {
            elseBlock = block();
        }
        CHECK_NULL_RETURN(elseBlock, nullptr);
        ifStmt->setElseBlock(elseBlock);
    }

    return ifStmt;
}

shared_ptr<ForStmt> StmtParser::forStmt() {
    CHECK_FALSE_RETURN(skipName(KeyWords::FOR), nullptr);
    CHECK_FALSE_RETURN(skipPunct("("), nullptr);

    auto initExpr = varDefStmt();
    CHECK_NULL_RETURN(initExpr, nullptr);

    auto conditionExpr = expr();
    CHECK_NULL_RETURN(initExpr, nullptr);
    CHECK_FALSE_RETURN(skipPunct(";"), nullptr);

    auto updateExpr = expr();
    CHECK_NULL_RETURN(updateExpr, nullptr);
    CHECK_FALSE_RETURN(skipPunct(")"), nullptr);

    auto block = this->block();
    CHECK_NULL_RETURN(block, nullptr);

    auto forStmt = make_shared<ForStmt>();
    forStmt->setInitExpr(initExpr);
    forStmt->setConditionExpr(conditionExpr);
    forStmt->setUpdateExpr(updateExpr);
    forStmt->setBlock(block);
    return forStmt;
}

shared_ptr<Ast> StmtParser::stmt() {
    auto token = tokenQueue.peek();
    if (token->type() == Lexer::TokenType::NAME) {
        auto nameToken = static_pointer_cast<Lexer::NameToken>(token);
        auto name = nameToken->getValue();
        if (name == KeyWords::VAR) {
            return varDefStmt();
        }
        if (name == KeyWords::IF) {
            return ifStmt();
        }
        if (name == KeyWords::FOR) {
            return forStmt();
        }
        if (name == KeyWords::RETURN) {
            return returnStmt();
        }
    }
    return simpleStmt();
}

/**
 * block: "{" {stmt} "}"
 */
shared_ptr<BlockStmt> StmtParser::block() {
    CHECK_FALSE_RETURN(skipPunct("{"), nullptr);

    vector<shared_ptr<Ast>> stmts;

    while (true) {
        if (isPunct(tokenQueue.peek(), "}")) {
            tokenQueue.poll();
            break;
        }

        auto stmt = this->stmt();
        CHECK_NULL_RETURN(stmt, nullptr);
        stmts.push_back(stmt);
    }

    return make_shared<BlockStmt>(stmts);
}

shared_ptr<ReturnStmt> StmtParser::returnStmt() {
    auto ret = skipName(KeyWords::RETURN);
    CHECK_FALSE_RETURN(ret, nullptr);

    if (isPunct(tokenQueue.peek(), ";")) {
        tokenQueue.poll();
        return make_shared<ReturnStmt>();
    }

    auto returnExpr = expr();
    CHECK_NULL_RETURN(returnExpr, nullptr);
    CHECK_FALSE_RETURN(skipPunct(";"), nullptr);

    return make_shared<ReturnStmt>(vector<shared_ptr<Ast>>{returnExpr});
}

shared_ptr<ImportStmt> StmtParser::importStmt() {
    auto ret = skipName(KeyWords::IMPORT);
    CHECK_FALSE_RETURN(ret, nullptr);

    auto name = this->name();
    CHECK_NULL_RETURN(name, nullptr);

    ret = skipPunct(";");
    CHECK_FALSE_RETURN(ret, nullptr);

    return make_shared<ImportStmt>(vector<shared_ptr<Ast>>{name});
}

} // namespace Parser
